from langchain_community.chat_message_histories import SQLChatMessageHistory

class LimitedSQLChatMessageHistory(SQLChatMessageHistory):
    def __init__(self, session_id: str, connection: str, max_records: int = 15):
        super().__init__(session_id, connection=connection)
        self.max_records = max_records


    def get_messages(self) -> list:
        data = super().get_messages()
        if len(data) > self.max_records:
            return [data[0]] + data[-self.max_records:]
        else:
            return data